import pandas as pd

def get_symbol(gdx, name, fields="L", columns=None, mappers=None):
    """ Get element from gdx as dataframe
    :param gdx: <gt.GdxContainer> with results
    :param name: <string> element name in gdx
    :param fields: <string or list> Values can be 'L', 'M', 'LO', 'UP', or 'SCALE' or 
            a list of any combination. Default is to format only 'L' (level) values 
            from GAMS into the final dict. Fields are ignored for GAMS sets.
    :param column: <list: string> column names 
    :param mappers: <dict: dict> {c_name: mapping}
            mapping applied to column c_name. 
            Mappings are applied after renameing columns
    """
    # skip loading if symbol already loaded in gdx
    if name not in gdx.__symLoaded__: 
        gdx.rgdx(name)
    df = gdx.to_dataframe(name, fields=fields)["elements"]
    if columns:
        df.columns = columns
    if mappers:
        for k,v in mappers.items():
            if v:
                df[k] = df[k].map(v)
    return df 

def get_date_mapping(gdx):
    """Get a mpping from date strings to dates"""
    df_date = get_symbol(gdx,"t")
    df_date = gdx.to_dataframe("t")["elements"]
    df_date["date"] = pd.to_datetime(df_date["t"])
    return df_date.set_index("t")["date"].to_dict()
    
def get_generation(gdx, map_dates=None):
    """Get generation variable from gdx file
    :param gdx: <gt.GdxContainer> with results
    :param map_dates: <dict> with mapping from date strings to datetime
    :return: <pd.DataFrame>"""
    df_gen = get_symbol(gdx, "GEN", 
                         columns=["plant", "date", "generation"],
                         mappers={"date": map_dates})
    # ensure values for all periods and each plant
    df_gen = df_gen.pivot_table("generation", "date", "plant"
                               ).fillna(0).stack().reset_index()
    df_gen.columns = ["date", "plant", "generation"]
    return df_gen

def get_elec_price(gdx, map_dates=None):
    """Get electricity price as marginal on market clearing constraint
    :param gdx: <gt.GdxContainer> with results
    :param map_dates: <dict> with mapping from date strings to datetime
    :return: <pd.DataFrame>"""    
    return get_symbol(gdx, "mkt_ele", fields="M",
                     columns=["date", "p_ele"],
                     mappers={"date": map_dates})

def get_capacity_rent(gdx, map_dates=None):
    """Get capacity rent as marginal on capapacity constraint
    :param gdx: <gt.GdxContainer> with results
    :param map_dates: <dict> with mapping from date strings to datetime
    :return: <pd.DataFrame>"""    
    cols = ["plant", "date", "p_cap"]
    df_pc = get_symbol(gdx, "mkt_cap", fields="M",
                      columns=["plant", "date", "p_cap"],
                      mappers={"date": map_dates})
    df_pc = df_pc.pivot_table("p_cap", "date", "plant"
                               ).fillna(0).stack().reset_index() 
    df_pc.columns = ["date", "plant", "p_cap"]
    return df_pc

def get_plant_characteristics(gdx):
    """Get plant characteristics from gdx file 
    :param gdx: <gt.GdxContainer> with results
    :return: <pd.DataFrame>
    """
    df = get_symbol(gdx, "r_characteristics", 
                    columns=["plant", "variable", "value"])
    df = df.pivot_table("value", "plant", "variable").reset_index()
    # merge fuel information
    df_fuel = get_symbol(gdx, "map_pf", columns=["plant", "fuel"])
    df = df.merge(df_fuel, on="plant")
    return df

def get_renewables(gdx, map_dates=None):
    """ Get renewable generation 
    :param gdx: <gt.GdxContainer> with results
    :param map_dates: <dict> with mapping from date strings to datetime
    :return: <pd.DataFrame>"""
    df = get_symbol(gdx, "ren",
                    columns=["renewable", "date", "value"],
                    mappers={"date": map_dates})
    return df.pivot_table("value", "date", "renewable").fillna(0).reset_index()

def get_cost(gdx, map_dates=None):
    """ Get marginal generation cost for each hour and plant
    :param gdx: <gt.GdxContainer> with results
    :param map_dates: <dict> with mapping from date strings to datetime
    :return: <pd.DataFrame>"""
    df = get_symbol(gdx, "r_cost",
                   columns=["plant", "date", "variable", "value"],
                   mappers={"date": map_dates})
    df = df.pivot_table("value", ["date", "plant"], "variable"
                       ).fillna(0).reset_index()
    # check that all cost catgories are represented
    for c in ["cost_total", "cost_carbon", "cost_fuel", "cost_ramping"]:
        if c not in df.columns:
            df[c] = 0		
    return df

def get_marginal_plant(df_gen, df_pc, df_eff, 
                       df_cost, df_pele):
    """ For each hour extract marginal plant
    :param df_gen: <pd.DataFrame> with hourly generation by plant
    :param df_pc: <pd.DataFrame> with hourly capacity rents by plant
    :param df_eff: <pd.DataFrame> with plant characteristics
    :param df_cost: <pd.DataFrame> with hourly marginal generation cost
    :param df_pele: <pd.DataFrame> with hourly electricitiy prices
    :return: <pd.DataFrame> with marignal plant and emission rate"""
    df_ = df_gen.merge(df_pc, on=["date", "plant"], how="outer"
                  ).fillna(0)
    # get frame with all active plants
    df_ = df_[(df_.generation > 1)].copy()
    # merge electricity prices and generation cost
    df_ = df_.merge(df_pele, on="date")
    df_ = df_.merge(df_cost, on=["date", "plant"])    
    # select the one with minimal markup over marginal generation costs
    df_["margin"] = df_.p_ele - (df_.cost_carbon + df_.cost_fuel)
    # df_ = df_.loc[df_.groupby("date").p_cap.idxmin()]  # select the one with minimal capacity rent
    df_ = df_.loc[df_.groupby("date").margin.idxmin()]

    return df_.merge(df_eff, on="plant", how="left")

def optimal_subsidy(df_marg, df_ren, external_cost=50,
                    periods=None):
    """ Calculate optimal subsidies for wind and solar power
    :param df_marg: <pd.DataFrame> with marginal plant
    :param df_ren: <pd.DataFrame> with renewable generation
    :param external_cost: <float> external cost of carbon [euro/tCO2]
    :param periods: <list: datetime> if provided data are 
            are filtered to only include values for respective periods
            Dates are assumed to be stored in "date" column
    """
    # merge all in one frame
    df = df_marg.merge(df_ren, on=["date"], how="outer")   
    
    # filter for given periods and convert re profiles to relative 
    if periods is None:
        periods = list(df_marg.date.unique())
    print(f"Get results for {len(periods)} of {len(df_marg)} periods")
    df = df[df.date.isin(periods)].copy()
    ren = ["solar", "wind"]
    for r in ren:
        df[r] = df[r]/df[r].sum()

    # marginal generation offsets
    df_off_gen = df.groupby("fuel")[["solar", "wind"]].sum()
    df_off_gen.index = [i + "_offset" for i in df_off_gen.index] 
    
    def weigted_average_renewables(col, col_new):
        """Small function taking the weighted average over column using
        renewable generation as weights. Its necessary due to somewhat unfortune
        data organization...
        Note that we use dataframes as well as list of renewable (ren) from
        outer scope
        :param col: <string> name of the column to aggregate
        :param col_new: <string> name of the new column
        :return: <pd.DataFrame>"""
        df_ = df.loc[:, ren + [col]]
        for r  in ren:
            df_[r] = df_[r]*df_[col]     
        return df_[ren].sum().to_frame(col_new)    
    
    # we have marginal generation offsets (Table 5) 
    # so lets calaculate optimal subsidy (components)
    # i.e. Table 6
    # emission offset
    df_off_emi = weigted_average_renewables("emission_rate", "marg_emission_offset").T
    # external benefit
    df_external = df_off_emi.copy()*external_cost
    df_external.index = ["external_benefit"]
    # market value
    df_market = weigted_average_renewables("p_ele", "market_value").T
    # avoided operating cost
    df_avoided = weigted_average_renewables("cost_fuel", "avoided_operating").T
    df_avoided_ramp = weigted_average_renewables("cost_ramping", "avoided_ramping").T 
    # carbon in market price
    df_val_carb = weigted_average_renewables("cost_carbon", "carbon_in_price").T
    # collect results
    lst_res = [df_off_gen, df_off_emi, df_external, 
               df_market, df_avoided, df_avoided_ramp,
               df_val_carb]
    df_res = pd.concat(lst_res)
    # optimal feedin
    df_feed = (df_res.loc["external_benefit",:] 
               + df_res.loc["avoided_operating",:]
               + df_res.loc["avoided_ramping",:]).to_frame("feedin").T
    # optimal premium
    df_prem = (df_res.loc["external_benefit",:] 
                   - df_res.loc["carbon_in_price",:]).to_frame("premium").T
    df_res = pd.concat([df_res, df_prem, df_feed])
    df_res.columns.name = "renewable"
    df_res.index.name = "variable"
    
    # finally lets get a tidy format
    df_res = df_res.stack().to_frame("value").reset_index()
    
    return df_res
	
def get_all_results(gdx):
    """Load all data and return optimal susbidy 
       over the whole sample period
    :param gdx: <gt.GdxContainer> with scenario results"""
    map_dates = get_date_mapping(gdx)
    df_gen = get_generation(gdx, map_dates)
    df_pele = get_elec_price(gdx, map_dates)
    df_pc = get_capacity_rent(gdx, map_dates)
    df_plants = get_plant_characteristics(gdx)
    df_ren = get_renewables(gdx, map_dates)
    df_cost = get_cost(gdx, map_dates)
    df_marg = get_marginal_plant(df_gen, df_pc, df_plants, df_cost, df_pele)
    df_opt = optimal_subsidy(df_marg, df_ren) 
    return df_opt